#include <string.h>
#include <assert.h>
#include "encrypt.h"

enum
{
    e_tea_loop = 16,
    e_tea_log_loop = 4,
    e_tea_delta = 0x9e3779b9,
};

/* src, dest: 8 bytes, key: 16 bytes */
static void encrypt_unit(const char* src, char* dest, const char* key)
{
    uint32_t sum = 0;
    uint32_t n = e_tea_loop;
    uint32_t src_0 = *(uint32_t*)&src[0];
    uint32_t src_1 = *(uint32_t*)&src[4];
    const uint32_t* key_int = (const uint32_t*)key;
    while (n-- > 0)
    {
        sum += e_tea_delta;
        src_0 += ((src_1 << 4) + key_int[0]) ^ (src_1 + sum) ^ ((src_1 >> 5) + key_int[1]);
        src_1 += ((src_0 << 4) + key_int[2]) ^ (src_0 + sum) ^ ((src_0 >> 5) + key_int[3]);
    }
 
    *(uint32_t*)(&dest[0]) = src_0;
    *(uint32_t*)(&dest[4]) = src_1;
}

/* src, dest: 8bytes, key: 16 bytes */
static void decrypt_unit(const char* src, char* dest, const char* key)
{
    uint32_t n = e_tea_loop;
    uint32_t sum = (e_tea_delta  << e_tea_log_loop);
    uint32_t src_0 = *(uint32_t*)&src[0];
    uint32_t src_1 = *(uint32_t*)&src[4];
    const uint32_t* key_int = (const uint32_t*)key;
    while (n-- > 0)
    {
        src_1 -= ((src_0 << 4) + key_int[2]) ^ (src_0 + sum) ^ ((src_0 >> 5) + key_int[3]);
        src_0 -= ((src_1 << 4) + key_int[0]) ^ (src_1 + sum) ^ ((src_1 >> 5) + key_int[1]);
        sum -= e_tea_delta;
    }
    *(uint32_t*)(&dest[0]) = src_0;
    *(uint32_t*)(&dest[4]) = src_1;
}


/* tea encrypt */
int32_t base_encrypt(const char* src, size_t src_len, int32_t key, char* dst, size_t* dst_len)
{
	assert(src && dst && dst_len);

	/* align by 8*/
	uint64_t align_add = 0;
	uint32_t src_floor = src_len / 8 * 8;
	uint32_t src_mod = src_len - src_floor;
	if(src_mod > 0)
	{
		memcpy((char*)&align_add, &src[src_floor], src_mod);
	}

	/* check dst len */
	if(*dst_len < src_mod + src_len + 8)
		return -1;

	/* encrypt: 8bytes head(pad) + tea encrypt */
	*(uint32_t*)dst = (src_mod > 0 ? 8 - src_mod : 0);

	/* gen 16 bytes key */
	char key_str[16];
	*(uint32_t*)key_str = key;
	*(uint32_t*)(key_str + 4) = (key ^ (key >> 1));
	*(uint32_t*)(key_str + 8) = (key ^ (key >> 2));
	*(uint32_t*)(key_str + 12) = (key ^ (key >> 3));
	
	int i;
	for(i=0; i<src_floor; i+=8)
		encrypt_unit(src + i, dst + i + 8, key_str);
	if(align_add > 0)
	{
		encrypt_unit((char*)&align_add, dst + src_floor + 8, key_str);
		*dst_len = src_floor + 8 + 8;
	}
	else
	{
		*dst_len = src_floor + 8;
	}
	
	return 0;
}

/* tea descrypt */
int32_t base_descrypt(const char* src, size_t src_len, int32_t key, char* dst, size_t* dst_len)
{
	assert(src && dst && dst_len);
	assert((src_len % 8 == 0) && (src_len > 8));
	
	if(*dst_len < src_len - 8)
		return -1;
	
	/* gen 16 bytes key */
	char key_str[16];
	*(uint32_t*)key_str = key;
	*(uint32_t*)(key_str + 4) = (key ^ (key >> 1));
	*(uint32_t*)(key_str + 8) = (key ^ (key >> 2));
	*(uint32_t*)(key_str + 12) = (key ^ (key >> 3));

	/* head */
	uint32_t pad = *(uint32_t*)src;

	/* descrypt by unit */
	int i;
	for(i=8; i<src_len; i+=8)
		decrypt_unit(src + i, dst + i - 8, key_str);
	*dst_len = src_len - 8 - pad;
	
	return 0;
}

